from .SSFM_each_span import SSFM_each_span
from .Amplifier import Amplifier
import numpy as np


def Fiber(sig, Signal, Fiber_para, PMD_para, Amplifer_para):
    # 更改数据格式
    sig = np.array(sig)
    sig = sig.transpose()
    sig = sig[0]
    for i in range(Fiber_para['Span']):
        Fiber_para["span_num"] = i
        # span_num = Fiber_para["span_num"]
        sig = SSFM_each_span(sig, Signal, Fiber_para, PMD_para)
        print("第", str(i+1), "Span分步傅里叶完成")  # 输出;将数值转换成字符串
        # 经过EDFA
        sig = Amplifier(sig, Amplifer_para, Signal, Fiber_para)
        print("第", str(i+1), "Span的Amplifer添加完成")
        signal = sig
        signal = np.array(signal)
        # single-pol
        signals_power = [np.mean(np.abs(signal) ** 2), np.mean(np.abs(signal) ** 2)]
        # dual-pol
        # signals_power = [np.mean(np.abs(signal[0:, 0]) ** 2), np.mean(np.abs(signal[0:, 1]) ** 2)]

    # 更改数据格式single-pol
    s1 = np.expand_dims(signal, axis=1)
    signal = s1

    # 更改数据格式dual-pol
    # s1 = np.expand_dims(signal[0:, 0], axis=1)
    # s2 = np.expand_dims(signal[0:, 1], axis=1)
    # signal = [s1, s2]

    return signal, signals_power
